# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from Account import hpQuant  # 回测模块
import Backtest
from Data import get_k_data  # 数据读取模块
import Setting as g  # 参数设置模块

#公式定义以及设定部分，引入所需参变量指标

def MA(Series, N):
    return pd.Series.rolling(Series, N).mean()


def CROSS(A, B):
    var = np.where(A < B, 1, 0)
    return (pd.Series(var, index=A.index).diff() < 0).apply(int)


def SMA(Series, N, M=1):
    ret = []
    i = 1
    length = len(Series)
    preY = Series.iloc[i]  # Y'
    ret.append(preY)
    while i < length:
        Y = (M * Series.iloc[i] + (N - M) * preY) / float(N)
        ret.append(Y)
        preY = Y
        i += 1
    return pd.Series(ret, index=Series.tail(len(ret)).index)


def REF(ser_, p):
    ser_, ref_ = list(ser_), []
    for i in range(p):
        ref_.append(np.nan)
    for i in range(len(ser_) - p):
        ref_.append(ser_[i])
    return ref_


def MAX(A, B):
    var = IF(A > B, A, B)
    return pd.Series(var, name='maxs')


def ABS(Series):
    return abs(Series)


def IF(COND, V1, V2):
    var = np.where(COND, V1, V2)
    return pd.Series(var, index=V1.index)

# 买卖策略设计与执行
# 首先要对数据预处理
df = get_k_data(ss=g.code, start=g.start_date, end=g.end_date)
mydf = df.copy()

C = mydf['close']
L = mydf['low']
H = mydf['high']
O = mydf['open']
V = mydf['volume']

ma5 = MA(C, 5)  # 生成均线 5日均线
mydf = mydf.join(pd.Series(ma5, name='MA5'))
ma10 = MA(C, 10)  # 生成均线 10日均线
mydf = mydf.join(pd.Series(ma10, name='MA10'))
b1 = CROSS(ma5, ma10)  # 5日均线上穿10日均线
s1 = CROSS(ma10, ma5)  # 5日均线下穿10日均线
mydf = mydf.join(pd.Series(s1, name='S1'))
mydf = mydf.join(pd.Series(b1, name='B1'))


# 第一部分回测：CTA策略的回测，调用class：hpQuant
tt = hpQuant()  # 初始化类
tt.code = g.code  # 证券代码，必须输入
tt.money = g.money
tt.money2 = g.money
tt.trading_Commission = g.trading_Commission
tt.stamp_duty = g.stamp_duty
tt.stop_loss_on = g.stop_loss_on
tt.priceStopLoss = g.priceStopLoss
tt.stop_loss_range = g.stop_loss_range
tt.stop_loss_max = g.stop_loss_max

# 参数表说明:股票数据表,买点序列名称,卖点序列名称,返回获利序列名称
# 输出第一部分回测结果
df3 = tt.Trade_testing(mydf, 'B1', 'S1', 'HL')  # 开始回测
print('\n打印交易过程\n')
print(repr(tt.trade_df))
print('\n 打印内部交易记录信息\n')
print(tt.text)

#第二部分：CTA+RSI策略组合的策略方法
def RSI(N1=6, N2=12, N3=24):
    """
    RSI 相对强弱指标分析
    """
    LC = REF(C, 1)
    RSI1 = SMA(MAX(C - LC, 0), N1, 1) / SMA(ABS(C - LC), N1, 1) * 100
    RSI2 = SMA(MAX(C - LC, 0), N2, 1) / SMA(ABS(C - LC), N2, 1) * 100
    RSI3 = SMA(MAX(C - LC, 0), N3, 1) / SMA(ABS(C - LC), N3, 1) * 100

    return RSI1, RSI2, RSI3


# 获取RSI指标
r1, r2, r3 = RSI(5, 10, 20)

mydf = mydf.join(pd.Series(r1, name='RSI1'))
mydf = mydf.join(pd.Series(r2, name='RSI2'))
mydf = mydf.join(pd.Series(r3, name='RSI3'))
mydf['S80'] = 80  # 增加上轨80轨迹线
mydf['X20'] = 20  # 增加下轨20轨迹线
mydf['Z50'] = 50  # 增加中轨50轨迹线


# 下面开始生成RSI指标买卖点
# 买点RSI1下穿20,且CTA不触发下交叉情况
b1 = CROSS(mydf['X20'], mydf['RSI1'])
mydf = mydf.join(pd.Series(b1, name='B1'))

# 卖点RSI1上穿80，且CTA不触发上交叉情况
s1 = CROSS(mydf['RSI1'], mydf['S80'])

# 卖点RSI1上穿50，且CTA不触发任何一种情况
s2 = CROSS(mydf['RSI1'], mydf['Z50'])

# 合并所有卖点信号
s3 = s1 | s2

mydf = mydf.join(pd.Series(s1, name='S1'))
mydf = mydf.join(pd.Series(s2, name='S2'))
mydf = mydf.join(pd.Series(s3, name='S3'))

# 第二部分回测调用
mm = hpQuant()  # 初始化类
mm.code = g.code  # 证券代码，必须输入
mm.money = g.money
mm.money2 = g.money
mm.trading_Commission = g.trading_Commission
mm.stamp_duty = g.stamp_duty
mm.stop_loss_on = g.stop_loss_on
mm.priceStopLoss = g.priceStopLoss
mm.stop_loss_range = g.stop_loss_range
mm.stop_loss_max = g.stop_loss_max

# 参数表说明:股票数据表,买点序列名称,卖点序列名称,返回获利序列名称
df3 = mm.Backtest.Trade_testing(mydf, 'B1', 'S3', 'HL')  # 开始回测
print('\n打印交易过程\n')
print(repr(mm.trade_df))
print('\n 打印内部交易记录信息\n')
print(mm.text)  # 打印交易信息

# 下面是绘图(第一部分单策略结果）
plt.figure(figsize=(12, 8), dpi=100)
ax1 = plt.subplot(211)
ax2 = plt.subplot(212)
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
# ax1（211窗口）
plt.sca(ax1)
plt.title(g.code + '  双均线回测演示')
# 显示网格：grid='on'
df3.close.plot(color='red', grid=True, legend=True)
df3['MA5'].plot(color='blue', grid=True, legend=True)
df3['MA10'].plot(color='green', grid=True, legend=True)
ax2.axhline(0, color='blue')

# ax2（212窗口）
plt.sca(ax2)
df3.HL.plot(color='orange', grid=True, legend=True)
df3.B1.plot(color='red', legend=True)
df3.S1.plot(color='blue', legend=True)
# 添加标题
plt.title(g.code + '  获利')
plt.show()

#第二部分，交叉策略结果
# 开启一个双图例的窗口，定义为211和212
plt.figure(figsize=(12, 8), dpi=100)
ax1 = plt.subplot(311)
ax2 = plt.subplot(312)
ax3 = plt.subplot(313)
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
# ax1（311窗口）
plt.sca(ax1)
plt.title(g.code + '  RSI指标买卖策略回测图')
# 显示网格：grid='on'
df3.close.plot(color='red', grid=True, legend=True)
df3['MA5'].plot(color='blue', grid=True, legend=True)
df3['MA10'].plot(color='green', grid=True, legend=True)
ax2.axhline(0, color='blue')

# ax2（312窗口）
plt.sca(ax2)
mydf.S80.plot.line()
mydf.X20.plot.line()
mydf.RSI1.plot.line(legend=True)
mydf.RSI2.plot.line(legend=True)
mydf.RSI3.plot.line(legend=True)

# ax3（313窗口）
plt.sca(ax3)
df3.HL.plot(color='orange', grid=True, legend=True)
df3.B1.plot(color='red', legend=True)
df3.S1.plot(color='blue', legend=True)
# 添加标题
plt.title(g.code + '  获利')
plt.show()
